{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Build a Multi-Agent Research Pipeline with LangGraph\n",
    "\n",
    "![Multi-Agent Research Pipeline](images/multi-agent-research-pipeline.png)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "llm = ChatOpenAI(model=\"gpt-4o-mini\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build a team of Analysts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "from typing_extensions import TypedDict\n",
    "from pydantic import BaseModel, Field\n",
    "\n",
    "class Analyst(BaseModel):\n",
    "    organization: str = Field(\n",
    "        description=\"Where the analyst works or is connected to.\",\n",
    "    )\n",
    "    full_name: str = Field(\n",
    "        description=\"The analyst's full name.\"\n",
    "    )\n",
    "    job_title: str = Field(\n",
    "        description=\"What the analyst does in this project or topic.\",\n",
    "    )\n",
    "    about: str = Field(\n",
    "        description=\"What the analyst cares about, what they focus on, or why they're involved.\",\n",
    "    )\n",
    "    @property\n",
    "    def profile(self) -> str:\n",
    "        return f\"\"\"\n",
    "Name: {self.full_name}\n",
    "Title: {self.job_title}\n",
    "Organization: {self.organization}\n",
    "About: {self.about}\n",
    "        \"\"\"\n",
    "\n",
    "class TeamOfAnalysts(BaseModel):\n",
    "    analysts: List[Analyst] = Field(\n",
    "        description=\"A list of people giving their input on the topic.\",\n",
    "    )\n",
    "\n",
    "class AnalystSetup(TypedDict):\n",
    "    topic: str                  # What this team is researching or discussing\n",
    "    number_of_analysts: int     # How many analysts we want to include\n",
    "    editor_feedback: str         # Notes or feedback from a human reviewer\n",
    "    analysts: List[Analyst]     # The current list of analysts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import Image, display\n",
    "from langgraph.graph import START, END, StateGraph\n",
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "from langchain_core.messages import AIMessage, HumanMessage, SystemMessage\n",
    "\n",
    "\n",
    "def build_team_of_analysts(state: AnalystSetup):\n",
    "    system_message = f\"\"\"\n",
    "    Your task is to create a team of AI analysts. Please follow these steps:\n",
    "\t\n",
    "    1.\tRead the research topic:\n",
    "    {state['topic']}\n",
    "\n",
    "\t2.\tLook at the optional notes or feedback from a human editor. It may help shape the analysts:\n",
    "    {state.get('editor_feedback', '')}\n",
    "\n",
    "\t3.\tFind the most interesting ideas or issues based on the topic and feedback.\n",
    "\t\n",
    "    4.\tChoose the top {state['number_of_analysts']} ideas.\n",
    "\t\n",
    "    5.\tCreate one analyst for each idea — each analyst should focus on just one theme.\n",
    "    \"\"\"\n",
    "    \n",
    "    structured_llm = llm.with_structured_output(TeamOfAnalysts)\n",
    "    analysts = structured_llm.invoke([SystemMessage(content=system_message), HumanMessage(content=\"Create a team of analysts for this topic.\")])\n",
    "    \n",
    "    return {\"analysts\": analysts.analysts}\n",
    "\n",
    "def human_feedback(state: AnalystSetup):\n",
    "    pass\n",
    "\n",
    "def should_continue(state: AnalystSetup):\n",
    "    human_feedback = state.get('editor_feedback', None)\n",
    "    if human_feedback:\n",
    "        return \"build_team_of_analysts\"\n",
    "    return END\n",
    "\n",
    "# Add nodes and edges \n",
    "builder = StateGraph(AnalystSetup)\n",
    "builder.add_node(\"build_team_of_analysts\", build_team_of_analysts)\n",
    "builder.add_node(\"human_feedback\", human_feedback)\n",
    "builder.add_edge(START, \"build_team_of_analysts\")\n",
    "builder.add_edge(\"build_team_of_analysts\", \"human_feedback\")\n",
    "builder.add_conditional_edges(\"human_feedback\", should_continue, [\"build_team_of_analysts\", END])\n",
    "\n",
    "# Compile\n",
    "memory = MemorySaver()\n",
    "graph = builder.compile(interrupt_before=['human_feedback'], checkpointer=memory)\n",
    "\n",
    "# View\n",
    "display(Image(graph.get_graph(xray=1).draw_mermaid_png()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "thread = {\"configurable\": {\"thread_id\": 3}}\n",
    "\n",
    "setup = graph.invoke({\n",
    "    \"topic\": \"What Model Context Protocol is and why it makes AI smarter\", \n",
    "    \"number_of_analysts\": 3}, \n",
    "    thread)\n",
    "\n",
    "for analyst in setup['analysts']:\n",
    "    print(analyst.profile)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get state and look at next node\n",
    "state = graph.get_state(thread)\n",
    "state.next"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "graph.update_state(thread, {\n",
    "    \"editor_feedback\": \"I think it would be helpful to add an analyst who focuses on how developers actually use MCP in real-world coding.\"}, \n",
    "    as_node=\"human_feedback\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Continue the graph execution\n",
    "for event in graph.stream(None, thread, stream_mode=\"values\"):\n",
    "    analysts = event.get('analysts', '')\n",
    "    if analysts:\n",
    "        for analyst in analysts:\n",
    "            print(analyst.profile)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "graph.update_state(thread, {\"editor_feedback\": None}, as_node=\"human_feedback\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "graph.stream(None, thread, stream_mode=\"values\")\n",
    "final_state = graph.get_state(thread)\n",
    "print(final_state.next)\n",
    "\n",
    "analysts = final_state.values.get('analysts')\n",
    "for analyst in analysts:\n",
    "    print(analyst.profile)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Analysts Interview Experts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "from typing_extensions import Annotated\n",
    "from langgraph.graph import MessagesState\n",
    "import operator\n",
    "\n",
    "\n",
    "# Stores everything that happens during the chat between one analyst and one expert\n",
    "class InterviewSession(MessagesState):\n",
    "    max_turns: int                          # How many times the analyst and expert can go back and forth\n",
    "    sources: Annotated[List, operator.add]  # All search results found during the chat (Tavely, Wiki, etc.)\n",
    "    analyst: Analyst                   # The analyst who is asking the questions\n",
    "    full_conversation: str                  # The full interview as plain text (what they talked about)\n",
    "    report_sections: List[str]              # What the analyst wrote based on the interview (can be 1 or more sections)\n",
    "                                            # This is what we’ll use later to build the full report"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Asking a Question"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ask_question(state: InterviewSession):\n",
    "    \"\"\"Generates a question from the analyst to the expert, using the analyst's profile.\"\"\"\n",
    "\n",
    "    system_msg = SystemMessage(content=f\"\"\"\n",
    "You are an AI analyst interviewing an expert about a specific topic.\n",
    "\n",
    "Your job is to ask clear and thoughtful questions to get helpful, surprising, and specific answers.\n",
    "\n",
    "1. Surprising: Ask things that lead to interesting or non-obvious insights.\n",
    "2. Specific: Avoid general talk — push for examples and real details.\n",
    "\n",
    "Here is your profile:\n",
    "{state[\"analyst\"].profile}\n",
    "\n",
    "Begin by introducing yourself in your analyst voice and ask your question.\n",
    "\n",
    "Keep asking until you understand the topic fully.\n",
    "\n",
    "When you're done, say: \"Thank you so much for your help!\" — this will end the interview.\n",
    "\n",
    "Stay in character throughout the conversation.\"\"\")\n",
    "\n",
    "    # Generate the question using the LLM\n",
    "    analyst_question = llm.invoke([system_msg] + state[\"messages\"])\n",
    "\n",
    "    # Return the new message to update the conversation\n",
    "    return {\"messages\": [analyst_question]}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generating an Answer\n",
    "\n",
    "#### Nodes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydantic import BaseModel, Field\n",
    "from langchain_core.messages import get_buffer_string\n",
    "from langchain_community.document_loaders import WikipediaLoader\n",
    "from langchain_community.tools import TavilySearchResults\n",
    "\n",
    "# A simple model to help the expert write a good search query\n",
    "class SearchTask(BaseModel):\n",
    "    search_text: str = Field(\n",
    "        None,\n",
    "        description=\"A short search query to help find useful information for the analyst’s question.\"\n",
    "    )\n",
    "\n",
    "search_prompt = SystemMessage(content=f\"\"\"\n",
    "You are helping generate a search query for a web search.\n",
    "\n",
    "You'll be given the full conversation between an analyst and an expert.  \n",
    "Look at the entire discussion, and focus especially on the **last question** from the analyst.\n",
    "\n",
    "Your task: Turn that question into a clear, well-structured search query.\"\"\")\n",
    "\n",
    "\n",
    "def search_web(state: InterviewSession):\n",
    "    \"\"\"Uses web search (Tavily) to find documents that help answer the analyst's question.\"\"\"\n",
    "\n",
    "    structured_llm = llm.with_structured_output(SearchTask)\n",
    "    search_query = structured_llm.invoke([search_prompt] + state[\"messages\"])\n",
    "\n",
    "    # Run Tavily search with the generated query\n",
    "    tavily_search = TavilySearchResults(max_results=5)\n",
    "    results = tavily_search.invoke(search_query.search_text)\n",
    "\n",
    "    # Format results into <Document> blocks\n",
    "    formatted_docs = \"\\n\\n---\\n\\n\".join(\n",
    "        [\n",
    "            f'<Document href=\"{doc[\"url\"]}\"/>\\n{doc[\"content\"]}\\n</Document>'\n",
    "            for doc in results\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    return {\"sources\": [formatted_docs]}\n",
    "\n",
    "\n",
    "def search_wikipedia(state: InterviewSession):\n",
    "    \"\"\"Uses Wikipedia to find documents that help answer the analyst's question.\"\"\"\n",
    "\n",
    "    structured_llm = llm.with_structured_output(SearchTask)\n",
    "    search_query = structured_llm.invoke([search_prompt] + state[\"messages\"])\n",
    "\n",
    "    # Run Wikipedia search using the query\n",
    "    results = WikipediaLoader(query=search_query.search_text, load_max_docs=5).load()\n",
    "\n",
    "    # Format results into <Document> blocks with source and page\n",
    "    formatted_docs = \"\\n\\n---\\n\\n\".join(\n",
    "        [\n",
    "            f'<Document source=\"{doc.metadata[\"source\"]}\" page=\"{doc.metadata.get(\"page\", \"\")}\"/>\\n{doc.page_content}\\n</Document>'\n",
    "            for doc in results\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    return {\"sources\": [formatted_docs]}\n",
    "\n",
    "\n",
    "def answer_question(state: InterviewSession):\n",
    "    \"\"\"Expert reads the question and answers it using only the documents found in search.\"\"\"\n",
    "\n",
    "    system_msg = SystemMessage(content=f\"\"\"\n",
    "You are an expert being interviewed by an AI analyst.\n",
    "\n",
    "Here is the analyst's profile:\n",
    "{state[\"analyst\"].profile}\n",
    "\n",
    "And here are documents you should use to answer the question:\n",
    "{state[\"sources\"]}\n",
    "\n",
    "Guidelines:\n",
    "1. Use only the info from the documents.\n",
    "2. Don't guess or add anything new.\n",
    "3. Reference documents using numbers like [1], [2].\n",
    "4. List those sources at the bottom.\n",
    "5. For example, write: [1] assistant/docs/mcp_guide.pdf, page 7.\"\"\")\n",
    "\n",
    "    expert_reply = llm.invoke([system_msg] + state[\"messages\"])\n",
    "    expert_reply.name = \"expert\"\n",
    "\n",
    "    return {\"messages\": [expert_reply]}\n",
    "\n",
    "\n",
    "def save_interview(state: InterviewSession):\n",
    "    \"\"\"Saves the full chat between analyst and expert as a plain text string.\"\"\"\n",
    "\n",
    "    conversation = get_buffer_string(state[\"messages\"])\n",
    "    return {\"full_conversation\": conversation}\n",
    "\n",
    "\n",
    "def write_report_section(state: InterviewSession):\n",
    "    \"\"\"\n",
    "    Writes a short report section based on the interview transcript,\n",
    "    supported by any referenced source documents.\n",
    "    \"\"\"\n",
    "    \n",
    "    system_msg = SystemMessage(content=f\"\"\"\n",
    "You are a technical writer creating a short report based on an interview with an expert.\n",
    "\n",
    "Your job is to write a clear, engaging section using the interview transcript as the main source, \n",
    "while using the attached documents only to support factual claims with proper citations.\n",
    "\n",
    "Here’s how to structure the report using Markdown:\n",
    "\n",
    "## Title  \n",
    "### Summary  \n",
    "### Sources\n",
    "\n",
    "Writing instructions:\n",
    "1. Use the interview transcript as your **primary source of insight**.\n",
    "2. If a factual claim in the interview **can be confirmed by a document**, cite it using [1], [2], etc.\n",
    "3. If a fact appears in the interview **but not in the documents**, it's okay to include it — just treat it as part of the expert's opinion.\n",
    "4. Do **not** invent or assume anything beyond the transcript and the documents.\n",
    "5. Keep the tone clear and concise. Avoid naming the expert or analyst.\n",
    "6. Max length: 500 words.\n",
    "7. In the “Sources” section, list each unique document used (no duplicates).\n",
    "8. Use full links or filenames (e.g., [1] https://example.com or assistant/docs/mcp_guide.pdf, page 7).\n",
    "\n",
    "Final review:\n",
    "- Ensure Markdown structure is followed\n",
    "- Make the title engaging and relevant to this focus area:\n",
    "  **{state[\"analyst\"].about}**\"\"\")\n",
    "\n",
    "    # Provide both the documents and interview to the LLM\n",
    "    user_msg = HumanMessage(content=f\"\"\"\n",
    "Here are the materials you should use:\n",
    "\n",
    "--- INTERVIEW TRANSCRIPT ---\n",
    "{state[\"full_conversation\"]}\n",
    "\n",
    "--- DOCUMENTS FOR CITATION ---\n",
    "{state[\"sources\"]}\n",
    "\"\"\")\n",
    "\n",
    "    report = llm.invoke([system_msg, user_msg])\n",
    "    return {\"report_sections\": [report.content]}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Conditional Edge"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def continue_or_finish(state: InterviewSession, name: str = \"expert\"):\n",
    "    \"\"\"Decides if the interview should continue or end after each answer.\"\"\"\n",
    "\n",
    "    messages = state[\"messages\"]\n",
    "    max_turns = state.get(\"max_turns\", 2)\n",
    "\n",
    "    # Count how many times the expert has responded\n",
    "    answers = [m for m in messages if isinstance(m, AIMessage) and m.name == name]\n",
    "    \n",
    "    if len(answers) >= max_turns:\n",
    "        return \"save_interview\"\n",
    "\n",
    "    # Check if the last analyst question said \"thank you\"\n",
    "    last_question = messages[-2]\n",
    "    if \"Thank you so much for your help\" in last_question.content:\n",
    "        return \"save_interview\"\n",
    "\n",
    "    return \"ask_question\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Constructing a Graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the graph and define its data structure\n",
    "interview_graph_builder = StateGraph(InterviewSession)\n",
    "\n",
    "# Define each step (node) in the conversation flow\n",
    "interview_graph_builder.add_node(\"ask_question\", ask_question)\n",
    "interview_graph_builder.add_node(\"search_web\", search_web)\n",
    "interview_graph_builder.add_node(\"search_wikipedia\", search_wikipedia)\n",
    "interview_graph_builder.add_node(\"answer_question\", answer_question)\n",
    "interview_graph_builder.add_node(\"save_interview\", save_interview)\n",
    "interview_graph_builder.add_node(\"write_report_section\", write_report_section)\n",
    "\n",
    "# Define the flow between steps\n",
    "\n",
    "# Start: analyst asks a question\n",
    "interview_graph_builder.add_edge(START, \"ask_question\")\n",
    "\n",
    "# Then search both web and Wikipedia in parallel\n",
    "interview_graph_builder.add_edge(\"ask_question\", \"search_web\")\n",
    "interview_graph_builder.add_edge(\"ask_question\", \"search_wikipedia\")\n",
    "\n",
    "# When both searches are done, expert answers the question\n",
    "interview_graph_builder.add_edge(\"search_web\", \"answer_question\")\n",
    "interview_graph_builder.add_edge(\"search_wikipedia\", \"answer_question\")\n",
    "\n",
    "# Decide whether to continue or end the interview\n",
    "interview_graph_builder.add_conditional_edges(\n",
    "    \"answer_question\",\n",
    "    continue_or_finish,\n",
    "    [\"ask_question\", \"save_interview\"]\n",
    ")\n",
    "\n",
    "# Once done, save the full interview transcript\n",
    "interview_graph_builder.add_edge(\"save_interview\", \"write_report_section\")\n",
    "\n",
    "# Finally, write the report section and finish\n",
    "interview_graph_builder.add_edge(\"write_report_section\", END)\n",
    "\n",
    "# Create the compiled graph with memory saving\n",
    "memory = MemorySaver()\n",
    "interview_graph = interview_graph_builder.compile(checkpointer=memory).with_config(\n",
    "    run_name=\"analysts_interview_experts\"\n",
    ")\n",
    "\n",
    "# Display the flow as an image\n",
    "display(Image(interview_graph.get_graph().draw_mermaid_png()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "analysts[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import Markdown\n",
    "\n",
    "thread = {\"configurable\": {\"thread_id\": \"1\"}}\n",
    "topic = \"What Model Context Protocol is and why it makes AI smarter\"\n",
    "\n",
    "# Start the conversation with a natural topic introduction\n",
    "messages = [HumanMessage(content=f\"I'm researching the topic: {topic}. I'd like to ask you a few questions.\")]\n",
    "\n",
    "\n",
    "interview = interview_graph.invoke({\n",
    "    \"analyst\": analysts[0],\n",
    "    \"messages\": messages,\n",
    "    \"max_turns\": 2\n",
    "}, thread)\n",
    "\n",
    "# Display the generated report section\n",
    "Markdown(interview[\"report_sections\"][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(interview[\"full_conversation\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build the Final Report"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List, Annotated\n",
    "import operator\n",
    "from typing_extensions import TypedDict\n",
    "\n",
    "class ResearchProjectState(TypedDict):\n",
    "    topic: str  # The main topic this project is focused on\n",
    "    number_of_analysts: int  # How many analysts we want to include\n",
    "    editor_feedback: str  # Notes or guidance from a human reviewer\n",
    "    analysts: List[Analyst]  # The team of analysts created for this topic\n",
    "\n",
    "    report_sections: Annotated[List[str], operator.add]  \n",
    "    # Sections written by each analyst based on their interview.\n",
    "    # These will be merged into a single list to build the final report.\n",
    "\n",
    "    introduction: str  # Intro paragraph for the final report\n",
    "    content: str       # Main content/body of the final report\n",
    "    conclusion: str    # Final paragraph that wraps up the report\n",
    "    final_report: str  # The full report (intro + content + conclusion)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Map - Reduce"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.constants import Send\n",
    "\n",
    "def start_all_interviews(state: ResearchProjectState):\n",
    "    \"\"\"\n",
    "    This is the map step of the graph.\n",
    "\n",
    "    If there's human feedback, we go back and rebuild the analyst team.\n",
    "\n",
    "    Otherwise, we run a separate interview sub-graph for each analyst,\n",
    "    passing in the topic as a message to start the conversation.\n",
    "    \"\"\"\n",
    "\n",
    "    # If feedback is present, return to team creation\n",
    "    editor_feedback = state.get(\"editor_feedback\")\n",
    "    if editor_feedback:\n",
    "        return \"build_team_of_analysts\"\n",
    "\n",
    "    # Otherwise, launch interviews in parallel\n",
    "    return [\n",
    "        Send(\n",
    "            \"analysts_interview_experts\",\n",
    "            {\n",
    "                \"analyst\": analyst,\n",
    "                \"messages\": [\n",
    "                    HumanMessage(content=f\"I'm researching the topic: {state[\"topic\"]}. I'd like to ask you a few questions.\")\n",
    "                ]\n",
    "            }\n",
    "        )\n",
    "        for analyst in state[\"analysts\"]\n",
    "    ]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create Writing Nodes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def write_report_summary(state: ResearchProjectState):\n",
    "    \"\"\"\n",
    "    Combines all analyst memos into the main report content.\n",
    "    \"\"\"\n",
    "\n",
    "    summary_instructions = f\"\"\"\n",
    "You are a technical writer. Your task is to write a report on the following topic:\n",
    "\n",
    "{state[\"topic\"]}\n",
    "\n",
    "Here's what you should know:\n",
    "\n",
    "- A team of analysts interviewed experts on different subtopics.\n",
    "- Each analyst wrote a short memo with their findings.\n",
    "\n",
    "Your job:\n",
    "\n",
    "1. You’ll receive all the memos written by the analysts.\n",
    "2. Carefully review them and extract the key ideas.\n",
    "3. Combine everything into one clear and cohesive summary that tells the full story.\n",
    "4. Organize the insights into a single, well-written report.\n",
    "\n",
    "Formatting instructions:\n",
    "\n",
    "- Use Markdown.\n",
    "- Don’t include any preamble or introduction to the report itself.\n",
    "- Don’t use any subheadings.\n",
    "- Begin with a top-level heading: ## Insights\n",
    "- Do not mention the names of any analysts.\n",
    "- Keep any citations from the memos (like [1], [2], etc.) exactly as they are.\n",
    "- At the end, add a ## Sources section with all sources used, listed in order.\n",
    "- Don’t repeat duplicate sources.\n",
    "\n",
    "Here’s how the sources section should look:\n",
    "\n",
    "[1] Source 1  \n",
    "[2] Source 2\n",
    "\n",
    "Here are the memos to work from:\n",
    "\n",
    "{\"\\n\\n\".join(state[\"report_sections\"])}\"\"\"\n",
    "\n",
    "\n",
    "    system_msg = SystemMessage(content=summary_instructions)\n",
    "    user_msg = HumanMessage(content=\"Please summarize all the memos into a single report.\")\n",
    "\n",
    "    summary = llm.invoke([system_msg, user_msg])\n",
    "    return {\"content\": summary.content}\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "def write_report_introduction(state: ResearchProjectState):\n",
    "    \"\"\"\n",
    "    Writes the introduction section for the final report.\n",
    "    \"\"\"\n",
    "\n",
    "    intro_instructions = f\"\"\"\n",
    "You are a technical writer helping to finalize a report on the topic:\n",
    "\n",
    "{state[\"topic\"]}\n",
    "\n",
    "You’ll receive all of the report’s main sections and will be asked to write either the introduction.\n",
    "\n",
    "Instructions:\n",
    "\n",
    "- Keep it clear and compelling.\n",
    "- Stick to around 100 words.\n",
    "- Don’t include any preamble before the section.\n",
    "- Use Markdown formatting.\n",
    "- Start with a strong title using a single # header\n",
    "- Then add a ## Introduction heading\n",
    "- Give a brief preview of the key ideas from the sections\n",
    "\n",
    "Here are the report sections to reflect on:\n",
    "\n",
    "{\"\\n\\n\".join(state[\"report_sections\"])}\"\"\"\n",
    "\n",
    "    \n",
    "    system_msg = SystemMessage(content=intro_instructions)\n",
    "    user_msg = HumanMessage(content=\"Please write an introduction for the report.\")\n",
    "\n",
    "    intro = llm.invoke([system_msg, user_msg])\n",
    "    return {\"introduction\": intro.content}\n",
    "\n",
    "\n",
    "def write_report_conclusion(state: ResearchProjectState):\n",
    "    \"\"\"\n",
    "    Writes the conclusion section for the final report.\n",
    "    \"\"\"\n",
    "\n",
    "    conclusion_instructions = f\"\"\"\n",
    "You are a technical writer helping to finalize a report on the topic:\n",
    "\n",
    "{state[\"topic\"]}\n",
    "\n",
    "You’ll receive all of the report’s main sections and will be asked to write either the conclusion.\n",
    "\n",
    "Instructions:\n",
    "\n",
    "- Keep it clear and compelling.\n",
    "- Stick to around 100 words.\n",
    "- Don’t include any preamble before the section.\n",
    "- Use Markdown formatting.\n",
    "- Use a ## Conclusion heading\n",
    "- Summarize the core takeaways from the report\n",
    "\n",
    "Here are the report sections to reflect on:\n",
    "\n",
    "{\"\\n\\n\".join(state[\"report_sections\"])}\"\"\"\n",
    "\n",
    "    system_msg = SystemMessage(content=conclusion_instructions)\n",
    "    user_msg = HumanMessage(content=\"Please write a conclusion for the report.\")\n",
    "\n",
    "    conclusion = llm.invoke([system_msg, user_msg])\n",
    "    return {\"conclusion\": conclusion.content}\n",
    "\n",
    "\n",
    "def finalize_report(state: ResearchProjectState):\n",
    "    \"\"\"\n",
    "    Combines the intro, main content, conclusion, and sources\n",
    "    into one complete Markdown-formatted report.\n",
    "    \"\"\"\n",
    "\n",
    "    content = state[\"content\"]\n",
    "    # Remove title if repeated inside the content\n",
    "    if content.startswith(\"## Insights\"):\n",
    "        content = content.replace(\"## Insights\", \"\").strip()\n",
    "\n",
    "    # Extract sources if present\n",
    "    if \"## Sources\" in content:\n",
    "        try:\n",
    "            content, sources = content.split(\"\\n## Sources\\n\")\n",
    "        except Exception:\n",
    "            sources = None\n",
    "    else:\n",
    "        sources = None\n",
    "\n",
    "    # Build final report\n",
    "    full_report = f\"\"\"\n",
    "    {state[\"introduction\"]}\n",
    "    \n",
    "    ---\n",
    "    \n",
    "    {content}\n",
    "    \n",
    "    ---\n",
    "    \n",
    "    {state[\"conclusion\"]}\"\"\"\n",
    "\n",
    "    if sources:\n",
    "        full_report += f\"\\n\\n## Sources\\n{sources}\"\n",
    "\n",
    "    return {\"final_report\": full_report}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Building Everything Up"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up the full research project graph\n",
    "builder = StateGraph(ResearchProjectState)\n",
    "\n",
    "# Add all nodes\n",
    "builder.add_node(\"build_team_of_analysts\", build_team_of_analysts)\n",
    "builder.add_node(\"human_feedback\", human_feedback)\n",
    "builder.add_node(\"analysts_interview_experts\", interview_graph)\n",
    "builder.add_node(\"write_report_summary\", write_report_summary)\n",
    "builder.add_node(\"write_report_introduction\", write_report_introduction)\n",
    "builder.add_node(\"write_report_conclusion\", write_report_conclusion)\n",
    "builder.add_node(\"finalize_report\", finalize_report)\n",
    "\n",
    "# Set up graph flow\n",
    "builder.add_edge(START, \"build_team_of_analysts\")\n",
    "builder.add_edge(\"build_team_of_analysts\", \"human_feedback\")\n",
    "builder.add_conditional_edges(\"human_feedback\", start_all_interviews, [\n",
    "    \"build_team_of_analysts\",\n",
    "    \"analysts_interview_experts\"\n",
    "])\n",
    "builder.add_edge(\"analysts_interview_experts\", \"write_report_summary\")\n",
    "builder.add_edge(\"analysts_interview_experts\", \"write_report_introduction\")\n",
    "builder.add_edge(\"analysts_interview_experts\", \"write_report_conclusion\")\n",
    "builder.add_edge(\n",
    "    [\"write_report_summary\", \"write_report_introduction\", \"write_report_conclusion\"],\n",
    "    \"finalize_report\"\n",
    ")\n",
    "builder.add_edge(\"finalize_report\", END)\n",
    "\n",
    "# Compile the graph\n",
    "memory = MemorySaver()\n",
    "research_graph = builder.compile(interrupt_before=[\"human_feedback\"], checkpointer=memory)\n",
    "\n",
    "# Display the full graph\n",
    "display(Image(research_graph.get_graph(xray=1).draw_mermaid_png()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Inputs\n",
    "thread = {\"configurable\": {\"thread_id\": 1}}\n",
    "number_of_analysts = 3 \n",
    "topic = \"What Model Context Protocol is and why it makes AI smarter\"\n",
    "\n",
    "\n",
    "# Run the graph until the first interruption\n",
    "for event in research_graph.stream({\"topic\": topic,\n",
    "                           \"number_of_analysts\": number_of_analysts}, \n",
    "                          thread, \n",
    "                          stream_mode=\"values\"):\n",
    "    analysts = event.get('analysts', '')\n",
    "    if analysts:\n",
    "        for analyst in analysts:\n",
    "            print(analyst.profile)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "state = research_graph.get_state(thread)\n",
    "state.next"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We now update the state as if we are the human_feedback node\n",
    "research_graph.update_state(thread, {\n",
    "    \"editor_feedback\": \"I think it would be helpful to add an analyst who focuses on how developers actually use MCP in real-world coding.\"\n",
    "    }, as_node=\"human_feedback\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check\n",
    "for event in research_graph.stream(None, thread, stream_mode=\"values\"):\n",
    "    analysts = event.get('analysts', '')\n",
    "    if analysts:\n",
    "        for analyst in analysts:\n",
    "            print(analyst.profile)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Confirm we are happy\n",
    "research_graph.update_state(thread, {\"editor_feedback\": None}, as_node=\"human_feedback\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Continue\n",
    "for event in research_graph.stream(None, thread, stream_mode=\"updates\"):\n",
    "    node_name = next(iter(event.keys()))\n",
    "    print(\"> \" + node_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import Markdown\n",
    "final_state = research_graph.get_state(thread)\n",
    "report = final_state.values.get('final_report')\n",
    "Markdown(report)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
